from playdoh03 import *
from base import *

import numpy
from numpy import inf, ones, zeros, tile, ones, histogram, nonzero, isscalar, maximum, minimum
from numpy.random import rand

class PSO(Optimization):
    @staticmethod
    def default_values(**opt_info):  
        if not 'omega' in opt_info.keys():
            opt_info['omega'] = .8
        if not 'cl' in opt_info.keys():
            opt_info['cl'] = 1.0
        if not 'cg' in opt_info.keys():
            opt_info['cg'] = 1.0
        if not 'Xmin' in opt_info.keys():
            opt_info['Xmin'] = -inf
        if not 'Xmax' in opt_info.keys():
            opt_info['Xmax'] = inf
        return opt_info
    
    def initialize(self):
        """
        Initializes the optimization algorithm
        Can use this variable :
            * self.opt_info = dict(algorithm-specific optimization parameters) # static
            self.this_machine_index
        """
        if self.this_machine_index == 0: # the master machine has index 0
            self.best_values = dict() # {machine: (best_position, best_fitness)}, filled by other machines
        
        self.D = len(self.params)
        self.N = self.particles
        self.V = zeros((self.D, self.N))
        
        Xmin = self.opt_info['Xmin']
        Xmax = self.opt_info['Xmax']
        if Xmin is None:
            Xmin = -inf*ones(self.D)
        if Xmax is None:
            Xmax = inf*ones(self.D)
        self.Xmin = Xmin
        self.Xmax = Xmax
        self.Xmin = tile(Xmin.reshape((-1,1)), (1, self.N))
        self.Xmax = tile(Xmax.reshape((-1,1)), (1, self.N))
        
        self.omega = self.opt_info['omega']
        self.cl = self.opt_info['cl']
        self.cg = self.opt_info['cg']
        
        # Preparation of the optimization algorithm
        self.fitness_lbest = -inf*ones(self.N)
        self.fitness_gbest = -inf
        self.X_lbest = self.X
        self.X_gbest = None

    def send_best(self, sender_index, best_position, best_fitness):
        """
        Call from other machines to the master machine
        """
        self.best_values[sender_index] = (best_position, best_fitness)
        for index in xrange(len(self.all_machines)):
            if index not in self.best_values.keys():
                log.debug("machine %d not in best value %s" % (index, str(self.best_values.keys())))
                self.set_this_ready(False)
                return False
        # here : all machines are in self.best_values : the global best fitness can
        # now be computed and the master machine is then ready!
        log.debug("ready now!")
        self.set_this_ready(True)
    
    def receive_best(self, X_gbest, fitness_gbest):
        """
        Call from the master to the slave : returns the global best
        """
        self.X_gbest, self.fitness_gbest = X_gbest, fitness_gbest
        self.set_this_ready(True) # this slave is ready because it has received the global best
    
    def extract_best(self, best_values):
        best_pos = None
        best_fit = -inf
        for (pos, fit) in best_values:
            if fit > best_fit:
                best_pos = pos
                best_fit = fit
        return (best_pos, best_fit)
    
    def get_this_best(self):
        """
        Returns the global best pos/fit on the current machine 
        """
        max_fitness = self.fitness.max()
        if max_fitness > self.fitness_gbest:
            index_gbest = nonzero(self.fitness == max_fitness)[0]
            if not(isscalar(index_gbest)): # if several best: take the first 
                index_gbest = index_gbest[0]
            self.X_gbest = self.X[:,index_gbest]
            self.fitness_gbest = max_fitness
        return self.fitness_gbest

    def get_this_local_best(self):
        indices_lbest = nonzero(self.fitness > self.fitness_lbest)[0]
        if (len(indices_lbest)>0):
            self.X_lbest[:,indices_lbest] = self.X[:,indices_lbest]
            self.fitness_lbest[indices_lbest] = self.fitness[indices_lbest]

    def iterate(self):
        """
        Performs the iteration step of the algorithm.
        The following variables are available :
            * self.param_values = initial parameters # current param values
            * self.X = corresponding initial matrix # current matrix
            * self.fitness = fitness values of initial parameters # current fitness
            * self.iteration = 0 # current iteration
        The variable self.results must be filled by this method.
        The following methods can be used :
            * self.contact : contacts other machines 
            * self.get_fitness(matrix) computes the fitness of the given matrix
            * self.this_machine, other_machines, all_machines in self.contact()
        """
        # reinitializes the best values
        if self.this_machine_index == 0: # the master machine has index 0
            self.best_values = dict() # {machine: (best_position, best_fitness)}, filled by other machines
        
        # 1. find best fitness
        # contact machine 0,  unless I am 0 : in this case I call send_best directly
        self.get_this_local_best()
        self.get_this_best()
        
        if self.all_machines is not None:
            if self.this_machine_index == 0:
                # master machine : wait for the slave machines to send their best fit
                self.best_values[0] = (self.X_gbest, self.fitness_gbest)
                self.wait()
                
                # now, compute the global best
                self.X_gbest, self.fitness_gbest = self.extract_best(self.best_values.values())
                
                # finally, connect to the slaves and give them the global best
                n = len(self.other_machines)
                self.contact(self.other_machines, 'receive_best', [self.X_gbest]*n, [self.fitness_gbest]*n)
            else:
                # other : send and wait
                # Multiple machine optimization: need to have the best fitness among all machines
                # sends the best fitness to the master
                
                self.contact(self.all_machines[0], 'send_best', [self.this_machine_index], [self.X_gbest], [self.fitness_gbest])
                
                # Waits for the master to return the global best
                self.wait()
        
        # 2. update matrix
        rl = rand(self.D, self.N)
        rg = rand(self.D, self.N)
        X_gbest_expanded = tile(self.X_gbest.reshape((-1,1)), (1, self.N))
        self.V = self.omega*self.V + self.cl*rl*(self.X_lbest-self.X) + self.cg*rg*(X_gbest_expanded-self.X)
        self.X = self.X + self.V
        
        self.X = maximum(self.X, self.Xmin)
        self.X = minimum(self.X, self.Xmax)
        
        # 3. compute new fitness
        self.get_fitness()
        
        self.results['best_fit'] = self.fitness_gbest
        self.results['best_pos'] = self.X_gbest

    @staticmethod
    def finalize(results_machines):
        return results_machines[0]